import torch
import torchvision.transforms as transforms
from PIL import Image

from demo import model
from model import NeuralNet


def main():
    # net = model(8, 2)
    net = NeuralNet(8, 2)
    net.load_state_dict(torch.load('net.pth'))

    x = torch.tensor([10,125,70,26,115,31.1,0.205,41],dtype=torch.float32)
    im = torch.unsqueeze(x, dim=0)  # [N, C, H, W]
    print(im)
    #
    with torch.no_grad():
        outputs = net(im)
        print(outputs)
        predict = torch.max(outputs, dim=1)[1].numpy()
        print(predict)


if __name__ == '__main__':
    main()
